from toolz import get

import numpy as np
import skimage.color as sc

from helm.dynamic.engine.metrics import Average


def get_y_channel(img):
    return sc.rgb2ycbcr(img)[:, :, 0]


def psnr(a1, a2):
    mse = np.mean((a1 - a2) ** 2)
    return 20 * np.log10(255.0 / np.sqrt(mse))


class PSNR(Average):

    def __init__(self, shave, name="psnr"):
        super().__init__(name, self.output_transform)
        self.shave = shave

    def output_transform(self, output):
        y_pred, y_true, batch_size = get(["y_pred", "y_true", "batch_size"], output)
        hr = y_true.numpy().transpose([0, 2, 3, 1]).astype(np.uint8)
        sr = y_pred.numpy().transpose([0, 2, 3, 1]).astype(np.uint8)
        hr = hr[:, self.shave:-self.shave, self.shave:-self.shave, :]
        sr = sr[:, self.shave:-self.shave, self.shave:-self.shave, :]
        hr = np.stack([get_y_channel(im) for im in hr])
        sr = np.stack([get_y_channel(im) for im in sr])
        return psnr(hr, sr), batch_size
